import tensorflow.compat.v1 as tf
import numpy as np
import matplotlib.pyplot as plt

ALPHA = 1e-1
ITERS = 4000

np.random.seed(777)
tf.random.set_random_seed(777)

x_data = np.float32([
    [0, 0],
    [0, 1],
    [1, 0],
    [1, 1]
])
y_data = np.float32([0, 1, 1, 0]).reshape(-1, 1)

# placeholder
x = tf.placeholder(tf.float32, [None, 2], name='ph_x')
y = tf.placeholder(tf.float32, [None, 1], name='ph_y')

# weight and bias
w1 =tf.Variable(tf.random.normal([2, 3]), dtype=tf.float32, name='w1')
b1 =tf.Variable(tf.random.normal([1, 3]), dtype=tf.float32, name='b1')
w2 =tf.Variable(tf.random.normal([3, 1]), dtype=tf.float32, name='w2')
b2 =tf.Variable(tf.random.normal([1, 1]), dtype=tf.float32, name='b2')

# forward
z1 = tf.add(tf.matmul(x, w1), b1, name='z1')  # mx2, 2x3 => mx3
a1 = tf.sigmoid(z1, name='a1')  # mx3
z2 = tf.add(tf.matmul(a1, w2), b2, name='z2')  # mx3, 3x1 => mx1
a2 = tf.sigmoid(z2, name='a2')  # mx1

# cost
j = tf.divide(
    tf.matmul(tf.transpose(y), tf.log(a2))
    +
    tf.matmul(tf.transpose(1 - y), tf.log(1 - a2)),
    - tf.cast(tf.shape(y)[0], tf.float32),
    name='j'
)

# backward propagation
# trunk
dz2 = tf.subtract(a2, y, name='dz2')  # mx1
da1 = tf.matmul(dz2, tf.transpose(w2), name='da1')  # mx1, 1x3 => mx3
dz1 = tf.multiply(da1, a1 * (1 - a1))  # mx3
# derivative of weight and bias
dw2 = tf.divide(
    tf.matmul(tf.transpose(a1), dz2),
    tf.cast(tf.shape(a1)[0], dtype=tf.float32),
    name='dw2'
)  # 3xm,mx1 => 3x1
db2 = tf.reduce_mean(dz2, axis=0, name='db2')  # 1x1
dw1 = tf.divide(
    tf.matmul(tf.transpose(x), dz1),
    tf.cast(tf.shape(x)[0], dtype=tf.float32),
    name='dw1'
)  # 2xm,mx3 => 2x3
db1 = tf.reduce_mean(dz1, axis=0, name='db1')  # 1x3

# update weight and bias
update = [
    tf.assign(w1, w1 - ALPHA * dw1),
    tf.assign(b1, b1 - ALPHA * db1),
    tf.assign(w2, w2 - ALPHA * dw2),
    tf.assign(b2, b2 - ALPHA * db2),
]

# metric
acc = tf.reduce_mean(
    tf.cast(
        tf.equal(a2 > 0.5, y > 0.5),
        tf.float32
    ),
    name='acc'
)

# train
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    cost_his = np.zeros(ITERS)
    GROUP = int(np.ceil(ITERS / 20))
    for i in range(ITERS):
        _, cost, accv = sess.run([update, j, acc], feed_dict={x: x_data, y: y_data})
        cost = cost[0][0]
        cost_his[i] = cost
        if i % GROUP == 0:
            print(f'#{i + 1}: cost = {cost}, acc = {accv}')
    if i % GROUP != 0:
        print(f'#{i + 1}: cost = {cost}, acc = {accv}')

    # cost function curve
    plt.plot(cost_his)
    plt.show()

    # evaluation
    pred = sess.run(a2, feed_dict={x: x_data, y: y_data})
    for p, label in zip(pred, y_data):
        print(f'target: {label} => prediction: {p} ({int(p > 0.5)})')

    # accuracy
    accv = sess.run(acc, feed_dict={x: x_data, y: y_data})
    print(f'Accuracy: {accv}')
